from attentivernn import AttentiveRNN from autoencoder import AutoEncoder from torch import nn class Generator(nn.Module): def __init__(self, repetition, blocks=3, layers=1, input_ch=3, out_ch=32, kernel_size=None, stride=1, padding=1, groups=1, dilation=1): super(Generator, self).__init__() if kernel_size is None: kernel_size = [3, 3] self.attentiveRNN = AttentiveRNN( repetition, blocks=blocks, layers=layers, input_ch=input_ch, out_ch=out_ch, kernel_size=None, stride=stride, padding=padding, groups=groups, dilation=dilation ) self.autoencoder = AutoEncoder() self.blocks = blocks self.layers = layers self.input_ch = input_ch self.out_ch = out_ch self.kernel_size = kernel_size self.stride = stride self.padding = padding self.groups = groups self.dilation = dilation self.sigmoid = nn.Sigmoid() def forward(self, x): x, attention_map = self.attentiveRNN(x) x = self.autoencoder(x * attention_map) return x if __name__ == "__main__": import torch from torchinfo import summary torch.set_default_tensor_type(torch.FloatTensor) generator = Generator(3, blocks=2) batch_size = 2 summary(generator, input_size=(batch_size, 3, 960,540))